import os
import argparse
import fitz  # PyMuPDF
from PIL import Image

def pdf_to_images(pdf_path, output_folder, dpi=200, fmt='jpeg'):
    """
    将单个PDF文件转换为图片并保存到指定文件夹。
    
    参数:
        pdf_path (str): PDF文件的路径
        output_folder (str): 输出图片的文件夹路径
        dpi (int): 图片分辨率，默认为200
        fmt (str): 输出图片的格式，默认为jpeg，可选png、jpg等
    """
    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    # 打开PDF文件
    try:
        pdf_document = fitz.open(pdf_path)
        print(f"正在转换: {pdf_path}, 总页数: {pdf_document.page_count}")
        
        # 遍历每一页
        for page_num in range(pdf_document.page_count):
            page = pdf_document.load_page(page_num)
            
            # 设置渲染分辨率
            pix = page.get_pixmap(dpi=dpi)
            
            # 将Pixmap转换为PIL Image
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            
            # 保存图片
            image_filename = f"page_{page_num + 1}.{fmt}"
            image_path = os.path.join(output_folder, image_filename)
            img.save(image_path, format=fmt)
            print(f"已保存: {image_path}")
        
        print(f"完成: {pdf_path}")
        pdf_document.close()
    except Exception as e:
        print(f"转换 {pdf_path} 时发生错误: {e}")

def batch_convert_pdfs(input_folder, output_folder, dpi=200, fmt='jpeg'):
    """
    批量转换指定文件夹中的所有PDF文件为图片。
    
    参数:
        input_folder (str): 包含PDF文件的输入文件夹路径
        output_folder (str): 输出图片的文件夹路径
        dpi (int): 图片分辨率，默认为200
        fmt (str): 输出图片的格式，默认为jpeg，可选png、jpg等
    """
    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    # 遍历输入文件夹中的所有文件
    for filename in os.listdir(input_folder):
        if filename.lower().endswith('.pdf'):
            pdf_path = os.path.join(input_folder, filename)
            pdf_name = os.path.splitext(filename)[0]
            pdf_output_folder = os.path.join(output_folder, pdf_name)
            
            # 转换单个PDF文件
            pdf_to_images(pdf_path, pdf_output_folder, dpi, fmt)

if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='批量将PDF文件转换为图片')
    parser.add_argument('input_folder', type=str, help='包含PDF文件的输入文件夹路径')
    parser.add_argument('output_folder', type=str, help='输出图片的文件夹路径')
    parser.add_argument('--dpi', type=int, default=200, help='图片分辨率，默认为200')
    parser.add_argument('--format', type=str, default='jpeg', help='输出图片的格式，默认为jpeg')
    
    args = parser.parse_args()
    
    # 调用批量转换函数
    batch_convert_pdfs(args.input_folder, args.output_folder, args.dpi, args.format)